import argparse
import sys
import os
import gzip
import mmap
import numpy as np
import time
from multiprocessing import Manager
import multiprocessing as mp

def generate_key(line):
    # 根据前两列生成key
    columns = line.split()[:2]
    return '||'.join(columns)

def process_file(output_file):
     # 检查输出文件是否存在
    try:
        with open(output_file, 'r') as f:
            existing_keys = set()
            for line in f:
                existing_keys.add(generate_key(line))
    except FileNotFoundError:
        existing_keys = set()
    return existing_keys

def replace_and_clean(string):
    # 将'], '替换成';'
    string = string.replace('],[', ';')
    # 删除空格和左方括号
    string = string.replace(' ', '').replace('[', '').replace(']', '')#.replace('\'','')
    return string

def read_feature(input,feature_Q,existing_keys,batch_size=1,qsize_limit=20):
    print('begin read feature', flush=True)
    fea_list=[]
    with open(input, 'r') as f:
        for line in f:
            words=line.strip().split('\t')
            key='||'.join(words[:2])
            if key in existing_keys:
                print('{} is in output'.format(words[0]), flush=True)
                continue
            while feature_Q.qsize()>=qsize_limit:
                time.sleep(20)
            fea_list.append(line)         
            if len(fea_list)>=batch_size:
                feature_Q.put(fea_list)
                fea_list=[]
    if len(fea_list)>0:
        feature_Q.put(fea_list)

    
def process_read(feature_Q,output_Q,seen_indexes,batch_size=1,time_wait=3,output_limit=10):
    fea_list=[]
    while True:
        if feature_Q.empty():
            time.sleep(time_wait)
        read_data = feature_Q.get()
        if read_data == "kill":
            feature_Q.put("kill")
            break
        for read in read_data:
            while output_Q.qsize()>=output_limit:
                time.sleep(20)
            words=read.strip().split('\t')
            if len(words)!=11:
                    print('words is less than 11 with {}'.format(words[0]), flush=True)
                    continue
            if int(words[10])<10:
                print('mapq is too low with {}'.format(words[0]), flush=True)
                continue
            key='||'.join(words[:2])
            if key in seen_indexes:
                print('{} is in output'.format(words[0]), flush=True)
                continue
            norm_signals=replace_and_clean(words[4]).split(";")
            #print(norm_signals[0])
            #print(norm_signals[1])
            norm_signals_new=[]
            for sig in norm_signals:
                norm_signals_new.append(np.round([float(y) for y in sig.split(",")], decimals=6))
            try:
                pred_dorado=np.round(np.array([float(x) for x in words[6].split(",")]), decimals=6)
            except:
                print(words[6])
                continue
            try:
                pred_deepsignal=np.round(np.array([float(x) for x in words[7].split(",")]), decimals=6)
            except:
                print(words[7])
                continue
            try:
                pred_mean=np.round(np.array([float(x) for x in words[8].split(",")]), decimals=6)
            except:
                print(words[8])
                continue
            pred_dorado_text=','.join([str(x) for x in pred_dorado])
            pred_deepsignal_text=','.join([str(x) for x in pred_deepsignal])
            pred_mean_text=','.join([str(x) for x in pred_mean])
            norm_signals_text=';'.join([",".join([str(y) for y in x]) for x in norm_signals_new])
            fea_str='\t'.join([words[0],words[1],words[2],words[3],norm_signals_text,words[5],pred_dorado_text,pred_deepsignal_text,pred_mean_text,words[9],words[10]])
                    
            if key not in seen_indexes:
                fea_list.append(fea_str)
                seen_indexes.add(key)
            if len(fea_list)>=batch_size:
                output_Q.put(fea_list)
                fea_list=[]
    if len(fea_list)>0:
        output_Q.put(fea_list)

def _write_featurestr(write_fp, featurestr_q,time_wait = 3):
    print('write_process-{} starts'.format(os.getpid()))
    output=0
    if os.path.exists(write_fp):
        with open(write_fp, 'a') as wf:
            while True:
                # during test, it's ok without the sleep(time_wait)
                if featurestr_q.empty():
                    time.sleep(time_wait)
                    continue
                features_str = featurestr_q.get()
                if features_str == "kill":
                    print('output line {}'.format(output))
                    print('write_process-{} finished'.format(os.getpid()))
                    break
                for one_features_str in features_str:
                    output+=1
                    wf.write(one_features_str + "\n")
                wf.flush()
    else:
        with open(write_fp, 'w') as wf:
            while True:
                # during test, it's ok without the sleep(time_wait)
                if featurestr_q.empty():
                    time.sleep(time_wait)
                    continue
                features_str = featurestr_q.get()
                if features_str == "kill":
                    print('output line {}'.format(output))
                    print('write_process-{} finished'.format(os.getpid()))
                    break
                for one_features_str in features_str:
                    output+=1
                    wf.write(one_features_str + "\n")
                wf.flush()

def parse_args():
    parser = argparse.ArgumentParser("")
    parser.add_argument("--input", type=str,required=True)
    parser.add_argument("--output", type=str,required=True)
    parser.add_argument("--nproc", "-p", type=int,required=True)
    return parser.parse_args()


def main():
    args=parse_args()
    #flag='w'
    #if os.path.exists(args.output):
    #    flag='a'
    existing_keys=process_file(args.output)
    seen_indexes=existing_keys
    print('build existing keys', flush=True)
    manager = mp.Manager()
    feature_Q = manager.Queue()
    output_Q = manager.Queue()
  
    pd=mp.Process(target=read_feature,args=(args.input,feature_Q,existing_keys),name="fea_reader")
    pd.daemon = True
    pd.start()
    ex_dp=args.nproc - 2
    ex_procs = []
    for i in range(ex_dp):
        pb = mp.Process(target=process_read, args=(feature_Q,output_Q,seen_indexes),
                          name="pb_reader")
            
        pb.daemon = True
        pb.start()
        ex_procs.append(pb)
    p_w = mp.Process(target=_write_featurestr, args=(args.output, output_Q),
                     name="writer")
    p_w.daemon = True
    p_w.start()
    pd.join()
    feature_Q.put("kill")
    for pb in ex_procs:
        pb.join()
    output_Q.put("kill")
    p_w.join()
    
        


if __name__ == '__main__':
    sys.exit(main())

